import torch
import sys
import numpy as np
#from utils import Config
import collections
from collections import OrderedDict
sys.path.append("/home/data/zhangyh22/ABINet/ABINet_Mindspore/")
from modules.model_abinet_iter import ABINetIterModel
# 通过PyTorch参数文件，打印PyTorch的参数文件里所有参数的参数名和shape，返回参数字典
def _get_model(config):
    import importlib
    names = config.model_name.split('.')
    module_name, class_name = '.'.join(names[:-1]), names[-1]
    cls = getattr(importlib.import_module(module_name), class_name)
    model = cls(config)
    #logging.info(model)
    return model


def pytorch_params_test(pth_file):
    #config = Config('/home/data4/zyh/ABINet/configs/train_abinet.yaml')
    #par_dict = _get_model(config)
    #par_dict.load_state_dict(torch.load(pth_file, map_location='cpu'),strict=False)
    par_dict = torch.load(pth_file, map_location='cpu')
    pt_params = {}
    for key,value in par_dict.items():
        #parameter = par_dict[key]
        if(key == 'model'):
            pt_parms = value
            for key1,value1 in value.items():
                #print(key1)
                #print(key1, value1.numpy().shape)


                pt_params[key1] = value1.numpy()
    return pt_parms  #单元测试用到
    #return pt_params   #转换权重用到


def pytorch_params_convert_weight(pth_file):
    #config = Config('/home/data4/zyh/ABINet/configs/train_abinet.yaml')
    #par_dict = _get_model(config)
    #par_dict.load_state_dict(torch.load(pth_file, map_location='cpu'),strict=False)
    par_dict = torch.load(pth_file, map_location='cpu')
    pt_params = {}
    for key,value in par_dict.items():
        #parameter = par_dict[key]
        if(key == 'model'):
            pt_parms = value
            for key1,value1 in value.items():
                #print(key1)
                #print(key1, value1.numpy().shape)
                # if(key1.endswith('self_attn.in_proj_weight')):
                #     value1 = torch.tensor(np.zeros((1536,512)))
                # if(key1.endswith('self_attn.in_proj_bias')):
                #     value1 = torch.tensor(np.zeros((1536)))
                # if(key1.endswith('linear2.bias')):
                #     value1 = torch.tensor(np.zeros((512)))
                # if(key1.endswith('linear2.weight')):
                #     value1 = torch.tensor(np.zeros((512,2048)))
                # if(key1.endswith('linear1.bias')):
                #     value1 = torch.tensor(np.zeros((2048)))
                # if(key1.endswith('linear1.weight')):
                #     value1 = torch.tensor(np.zeros((2048,512)))
                # if(key1.endswith('self_attn.out_proj.weight')):
                #     value1 = torch.tensor(np.zeros((512,512)))
                # if(key1.endswith('self_attn.out_proj.bias') or key1.endswith('norm1.weight') or key1.endswith('norm1.bias') or key1.endswith('norm2.weight') or key1.endswith('norm2.bias')):
                #     value1 = torch.tensor(np.zeros((512)))                    
                pt_params[key1] = value1.numpy()
    #return pt_parms  #单元测试用到
    return pt_params   #转换权重用到

# 通过MindSpore的Cell，打印Cell里所有参数的参数名和shape，返回参数字典
def mindspore_params(network):
    #print(network)
    ms_params = {}
    for param in network.get_parameters():
        name = param.name
        value = param.data.asnumpy()
        #print(name)
        print(name, value.shape)
        ms_params[name] = value
    return ms_params
